#include <chrono>
#include <iostream>

#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>
#include <pcl/common/transforms.h>
#include <pcl/filters/approximate_voxel_grid.h>
#include <pcl/filters/voxel_grid.h>

#include "test_tools.hpp"
#include "pclview_tools.hpp"

int main(int argc, char** argv) {
    std::string config_file = "../config/config.yaml";
    // std::string config_file = "/home/qianliyue/rosdev/lio_learn_ws/src/testICP/config/config.yaml";
    YAML::Node config_node = YAML::LoadFile(config_file);
    std::string t_str = config_node["target_file"].as<std::string>();
    std::string s_str = config_node["source_file"].as<std::string>();
    std::cout << "load file: " << t_str << ", and " << s_str << std::endl;
    double ds_size = config_node["ds_size"].as<double>();

    CloudPtr cloud_source(new CloudType());
    CloudPtr cloud_target(new CloudType());
    if (pcl::io::loadPCDFile<PointType>(s_str, *cloud_source) == -1) {
        PCL_ERROR("Couldn't read file test_pcd.pcd\n");
        return (-1);
    }
    if (pcl::io::loadPCDFile<PointType>(t_str, *cloud_target) == -1) {
        PCL_ERROR("Couldn't read file test_pcd.pcd\n");
        return (-1);
    }

    pcl::VoxelGrid<PointType> ds_sample;
    ds_sample.setLeafSize(ds_size, ds_size, ds_size);
    ds_sample.setInputCloud(cloud_source);
    ds_sample.filter(*cloud_source);
    ds_sample.setInputCloud(cloud_target);
    ds_sample.filter(*cloud_target);

    std::cout << "target:" << cloud_target->size() << "[pts] source:" << cloud_source->size() << "[pts]" << std::endl;

// #define ALGO_NUM 12

    AlignResult* Res[11];

    /////////////////////////////////////////////////////////////////////////

    std::cout << "--- pcl_icp ---" << std::endl;
    pcl::IterativeClosestPoint<PointType, PointType> pcl_icp;
    Res[0] = test_pcl(pcl_icp, cloud_target, cloud_source);

    std::cout << "--- pcl_ndt ---" << std::endl;
    pcl::NormalDistributionsTransform<PointType, PointType> pcl_ndt;
    pcl_ndt.setResolution(1.0);
    Res[1] = test_pcl(pcl_ndt, cloud_target, cloud_source);

    std::cout << "--- pcl_gicp ---" << std::endl;
    pcl::GeneralizedIterativeClosestPoint<PointType, PointType> pcl_gicp;
    Res[2] = test_pcl(pcl_gicp, cloud_target, cloud_source);

    /////////////////////////////////////////////////////////////////////////

    std::cout << "--- fgicp_st ---" << std::endl;
    fast_gicp::FastGICPSingleThread<PointType, PointType> fgicp_st;
    test_fast(fgicp_st, cloud_target, cloud_source);

    std::cout << "--- fgicp_mt ---" << std::endl;
    fast_gicp::FastGICP<PointType, PointType> fgicp_mt;
    // fast_gicp uses all the CPU cores by default
    // fgicp_mt.setNumThreads(8);
    Res[3] = test_fast(fgicp_mt, cloud_target, cloud_source);

    std::cout << "--- vgicp_st ---" << std::endl;
    fast_gicp::FastVGICP<PointType, PointType> vgicp;
    vgicp.setResolution(1.0);
    vgicp.setNumThreads(1);
    test_fast(vgicp, cloud_target, cloud_source);

    std::cout << "--- vgicp_mt ---" << std::endl;
    vgicp.setNumThreads(omp_get_max_threads());
    Res[4] = test_fast(vgicp, cloud_target, cloud_source);

    /////////////////////////////////////////////////////////////////////////

    std::cout << "--- ICP_SVD ---" << std::endl;
    TESTICP::ICPSVDRegistration icp_svd(config_node["ICP_SVD"]);
    Res[5] = test_pcl(icp_svd, cloud_target, cloud_source);

    // std::cout << "--- ICP_GN ---" << std::endl;
    // TESTICP::opt_ICP_GN icp_gn(config_node["ICP_GN"]);
    // Res[6] = test_pcl(icp_gn, cloud_target, cloud_source);

    std::cout << "--- ICP_CERES ---" << std::endl;
    TESTICP::opt_ICP_CERES icp_ceres(config_node["ICP_CERES"]);
    Res[6] = test_pcl(icp_ceres, cloud_target, cloud_source);

    std::cout << "--- ICP_G2O ---" << std::endl;
    TESTICP::opt_ICP_G2O icp_g2o(config_node["ICP_G2O"]);
    Res[7] = test_pcl(icp_g2o, cloud_target, cloud_source);

    std::cout << "--- ICP_GTSAM ---" << std::endl;
    TESTICP::opt_ICP_GTSAM icp_gtsam(config_node["ICP_GTSAM"]);
    Res[8] = test_pcl(icp_gtsam, cloud_target, cloud_source);

    /////////////////////////////////////////////////////////////////////////

    std::cout << "--- KCP ---" << std::endl;
    Res[9] = test_kcp(cloud_target, cloud_source);


    /////////////////////////////////////////////////////////////////////////


    std::string fun_name[] = { "pcl_icp", "pcl_ndt", "pcl_gicp", "fgicp", "vgicp","ICP_SVD","ICP_CERES","ICP_G2O","ICP_GTSAM","KCP", "ori" };
    size_t ALGO_NUM = sizeof(fun_name)/sizeof(fun_name[0]);

    std::vector<CloudPtr> output_cloud(ALGO_NUM - 1);
    PclViewer<PointType> PclView(4, 3);

    for (int i = 0; i < ALGO_NUM - 1; i++) {
        output_cloud[i] = CloudPtr(new CloudType);
        pcl::transformPointCloud(*cloud_source, *output_cloud[i], Res[i]->tf);

        char text[64];
        std::snprintf(text, sizeof(text), "%s time:%.0fms score:%.4f", fun_name[i].c_str(), Res[i]->time, Res[i]->score);

        PclView.addTitle(i + 1, std::string(text));
        PclView.addCloud(output_cloud[i], "output_cloud", i + 1, { 0,255,0 });
        PclView.addCloud(cloud_target, "cloud_target", i + 1, { 255,0,0 });
    }
    PclView.addTitle(ALGO_NUM, fun_name[ALGO_NUM - 1]);
    PclView.addCloud(cloud_source, "output_cloud", ALGO_NUM, { 0,255,0 });
    PclView.addCloud(cloud_target, "cloud_target", ALGO_NUM, { 255,0,0 });

    PclView.show();

#ifdef USE_VGICP_CUDA
    std::cout << "--- ndt_cuda (P2D) ---" << std::endl;
    fast_gicp::NDTCuda<PointType, PointType> ndt_cuda;
    ndt_cuda.setResolution(1.0);
    ndt_cuda.setDistanceMode(fast_gicp::NDTDistanceMode::P2D);
    test(ndt_cuda, cloud_target, cloud_source);

    std::cout << "--- ndt_cuda (D2D) ---" << std::endl;
    ndt_cuda.setDistanceMode(fast_gicp::NDTDistanceMode::D2D);
    test(ndt_cuda, cloud_target, cloud_source);

    std::cout << "--- vgicp_cuda (parallel_kdtree) ---" << std::endl;
    fast_gicp::FastVGICPCuda<PointType, PointType> vgicp_cuda;
    vgicp_cuda.setResolution(1.0);
    // vgicp_cuda uses CPU-based parallel KDTree in covariance estimation by default
    // on a modern CPU, it is faster than GPU_BRUTEFORCE
    // vgicp_cuda.setNearestNeighborSearchMethod(fast_gicp::NearestNeighborMethod::CPU_PARALLEL_KDTREE);
    test(vgicp_cuda, cloud_target, cloud_source);

    std::cout << "--- vgicp_cuda (gpu_bruteforce) ---" << std::endl;
    // use GPU-based bruteforce nearest neighbor search for covariance estimation
    // this would be a good choice if your PC has a weak CPU and a strong GPU (e.g., NVIDIA Jetson)
    vgicp_cuda.setNearestNeighborSearchMethod(fast_gicp::NearestNeighborMethod::GPU_BRUTEFORCE);
    test(vgicp_cuda, cloud_target, cloud_source);

    std::cout << "--- vgicp_cuda (gpu_rbf_kernel) ---" << std::endl;
    // use RBF-kernel-based covariance estimation
    // extremely fast but maybe a bit inaccurate
    vgicp_cuda.setNearestNeighborSearchMethod(fast_gicp::NearestNeighborMethod::GPU_RBF_KERNEL);
    // kernel width (and distance threshold) need to be tuned
    vgicp_cuda.setKernelWidth(0.5);
    test(vgicp_cuda, cloud_target, cloud_source);
#endif

    return 0;
}
